import math

import paddle.fluid.layers as L

from helm.static.models.layers import Conv2d, DEFAULTS, GlobalAvgPool, Pool2d, Linear, Layer, Sequential


class Shortcut(Sequential):

    def __init__(self, in_channels, out_channels, stride):
        layers = []
        if in_channels != out_channels or stride != 1:
            if stride != 1:
                layers.append(Pool2d(2, 2, type='avg'))
            layers.append(Conv2d(in_channels, out_channels, 1, stride=1))
        super().__init__(*layers)


class Bottleneck(Layer):
    expansion = 4

    def __init__(self, in_channels, out_channels, stride, groups, base_width):
        super().__init__()

        D = math.floor(out_channels // self.expansion * (base_width / 64))

        self.conv1 = Conv2d(in_channels, D * groups, kernel_size=1, bn=True, act='default')
        self.conv2 = Conv2d(D * groups, D * groups, kernel_size=3, stride=stride, groups=groups,
                            bn=True, act='default')
        self.conv3 = Conv2d(D * groups, out_channels, kernel_size=1, bn=True)
        self.shortcut = Shortcut(in_channels, out_channels, stride)

    def forward(self, x):
        identity = self.shortcut(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = L.elementwise_add(x, identity, act=DEFAULTS['activation'])
        return x


class ResNeXt(Layer):

    def __init__(self, stem_channels=64, channels_per_stage=(256, 512, 1024), depth=29, groups=8, base_width=64,
                 num_classes=10):
        super().__init__()
        layers = [(depth - 2) // 9] * 3

        cs = [stem_channels] + list(channels_per_stage)

        self.conv = Conv2d(3, stem_channels, kernel_size=3)

        self.layer1 = self._make_layer(
            cs[0], cs[1], layers[0], stride=1, groups=groups, base_width=base_width)
        self.layer2 = self._make_layer(
            cs[1], cs[2], layers[1], stride=2, groups=groups, base_width=base_width)
        self.layer3 = self._make_layer(
            cs[2], cs[3], layers[2], stride=2, groups=groups, base_width=base_width)

        self.avgpool = GlobalAvgPool()
        self.fc = Linear(cs[3], num_classes)

    def _make_layer(self, in_channels, out_channels, blocks, stride, groups, base_width):
        layers = [Bottleneck(in_channels, out_channels, stride=stride, groups=groups, base_width=base_width)]
        for i in range(1, blocks):
            layers.append(
                Bottleneck(out_channels, out_channels, stride=1, groups=groups, base_width=base_width))
        return Sequential(*layers)

    def forward(self, x):
        x = self.conv(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x = self.avgpool(x)
        x = self.fc(x)
        return x


def test_resnext():

    x = L.randn((1, 3, 32, 32))
    #
    # # ResNeXt-29, 8×64d
    # net = ResNeXt(stages=(64, 256, 512, 1024), depth=29, groups=8, base_width=64)
    # assert profile(net, (x,))[1] == 34426634
    #
    # # ResNeXt-29, 16×64d
    # net = ResNeXt(stages=(64, 256, 512, 1024), depth=29, groups=16, base_width=64)
    # assert profile(net, (x,))[1] == 68155146
